import argparse
import pathlib

import cv2
import matplotlib.pyplot
from PIL import Image
from pylab import *

matplotlib.use('TkAgg')
import csv
import codecs


class PolyMasker:

    def __init__(self, n_point, img_file):
        self.n_point = n_point
        self.img_file = img_file
        p = pathlib.Path(img_file)
        self.img_mask_name = p.stem + "_mask" + p.suffix
        self.points = []
        self.img = None
        self.mask = None

    def run(self):
        self.img = cv2.imread(self.img_file)
        ion()
        imshow(self.img)
        self.points = ginput(self.n_point)
        ioff()
        show()
        self._draw_poly()

    def _draw_poly(self):
        self.mask = np.zeros(self.img.shape, dtype=np.uint8)
        pts = np.array(self.points, np.int32)  # 数据类型必须为 int32
        pts = pts.reshape((1, -1, 2))
        cv2.fillPoly(self.mask, [pts], color=(255, 255, 255))
        cv2.imwrite(self.img_mask_name, self.mask)


def parse_opt():
    parser = argparse.ArgumentParser()
    parser.add_argument('--file', type=str, default='walk_screenshot.png', help='保存的截图名称')
    parser.add_argument('--n', type=int, default=4, help='多边形点数')
    return parser.parse_args()


if __name__ == '__main__':
    opt = parse_opt()
    msk = PolyMasker(opt.n, opt.file)
    msk.run()